# 从封装ST_GCN模块下的train子模块中的train_model文件导入STGCNTrainer类
# 该类负责模型的训练过程
from 封装ST_GCN.train.train_model import STGCNTrainer
# 从封装ST_GCN模块下的test子模块中的test_model文件导入STGCNTester类
# 该类负责模型的测试过程
from 封装ST_GCN.test.test_model import STGCNTester


if __name__ == "__main__":
    # 定义训练数据文件路径
    train_data_path = ''
    # 定义训练标签文件路径
    train_label_path = ''
    # 定义测试数据文件路径
    test_data_path = ''
    # 定义测试标签文件路径
    test_label_path = ''

    # 定义数据集的人体节点数
    num_nodes = 25
    # 定义是否使用残差连接，True表示使用
    use_residual = True
    # 定义训练和测试时的批次大小
    batch_size = 32
    # 定义训练的总轮数
    num_epochs = 300
    # 定义训练时的学习率
    lr = 0.001
    # 定义早停策略的耐心值，即多少轮没有提升就停止训练
    patience = 30
    # 定义图的相关参数，使用openpose布局和adaptive策略
    graph_args = {'layout': 'openpose', 'strategy': 'adaptive'}
    # 定义是否使用边重要性加权，True表示使用
    edge_importance_weighting = True

    # 创建STGCNTrainer类的实例，传入训练数据路径、训练标签路径、节点数等参数
    trainer = STGCNTrainer(train_data_path, train_label_path, num_nodes, use_residual, batch_size, num_epochs, lr,
                           patience,
                           graph_args, edge_importance_weighting)
    # 调用trainer实例的train方法开始训练模型
    trainer.train()

    # 创建STGCNTester类的实例，传入测试数据路径、测试标签路径、节点数等参数
    tester = STGCNTester(test_data_path, test_label_path, num_nodes, batch_size, graph_args, edge_importance_weighting,
                         use_residual=use_residual)
    # 调用tester实例的load_model方法加载训练好的模型
    tester.load_model()
    # 调用tester实例的test方法对加载的模型进行测试
    tester.test()